import rospy
from geometry_msgs.msg import Twist
import numpy as np
from scipy.fft import fft
import threading
import time
from queue import Queue, Full, Empty
import sounddevice as sd

# 音频参数 - 使用更小的 CHUNK
RATE = 44100
CHUNK = 1024  # 减小块大小
QUEUE_SIZE = 10  # 增加队列大小

# 初始化ROS节点
rospy.init_node("sound_receiver_node")
cmd_pub = rospy.Publisher("/cmd_vel", Twist, queue_size=10)

# 音频队列
audio_queue = Queue(maxsize=QUEUE_SIZE)

# 频率映射表
key_to_freq = {
    "7": 18500,
    "8": 18700,
    "9": 18900,
    "4": 19100,
    "6": 18600,
    "1": 18800,
    "2": 19000,
    "3": 19200,
}

key_to_cmd = {
    "7": [1, 1],
    "8": [1, 0],
    "9": [1, -1],
    "4": [0, 1],
    "6": [0, -1],
    "1": [-1, 1],
    "2": [-1, 0],
    "3": [-1, -1],
}

freq_to_cmd = {
    f: {"x": c[0], "y": c[1]} for f, c in zip(key_to_freq.values(), key_to_cmd.values())
}


def sound_recognize(wave):
    """优化版频率识别算法"""
    n = len(wave)
    if n == 0:
        return None

    # 应用汉宁窗
    window = np.hanning(n)
    wave_windowed = wave * window

    # FFT计算
    try:
        fft_result = fft(wave_windowed)
        freqs = np.fft.fftfreq(n, 1 / RATE)

        # 只取正频率部分
        half_n = n // 2
        fft_magnitude = np.abs(fft_result[:half_n])
        freqs = freqs[:half_n]

        # 设置能量阈值
        if len(fft_magnitude) > 0:
            energy_threshold = np.mean(fft_magnitude) * 10
        else:
            return None

        best_freq = None
        max_energy = 0

        # 只检查目标频率附近的窄带
        for target_freq in freq_to_cmd.keys():
            # 计算目标频率在频谱中的位置
            idx = np.argmin(np.abs(freqs - target_freq))

            # 检查窄带范围 (目标频率±50Hz)
            start_idx = max(0, idx - int(50 / (RATE / n)))
            end_idx = min(half_n, idx + int(50 / (RATE / n)))

            # 计算窄带内的能量
            if start_idx < end_idx:
                energy = np.sum(fft_magnitude[start_idx:end_idx])
            else:
                energy = 0

            # 只有能量超过阈值才考虑
            if energy > energy_threshold and energy > max_energy:
                max_energy = energy
                best_freq = target_freq

        return best_freq
    except Exception as e:
        rospy.logwarn(f"频率识别错误: {str(e)}")
        return None


def audio_producer():
    """音频采集线程 - 使用 PulseAudio 并添加错误处理"""
    rospy.loginfo("音频采集已启动")

    # 获取所有音频设备
    try:
        devices = sd.query_devices()
    except:
        devices = []
        rospy.logwarn("无法获取音频设备列表")

    # 打印可用设备信息
    if devices:
        rospy.loginfo("可用音频输入设备:")
        for i, dev in enumerate(devices):
            if dev.get("max_input_channels", 0) > 0:
                samplerate = dev.get("default_samplerate", "未知")
                rospy.loginfo(
                    f"{i}: {dev['name']} (输入通道: {dev.get('max_input_channels', 0)}, 采样率: {samplerate})"
                )

    # 尝试使用 PulseAudio
    input_device = None
    for i, dev in enumerate(devices):
        if "pulse" in dev["name"].lower() and dev.get("max_input_channels", 0) > 0:
            input_device = i
            rospy.loginfo(f"使用 PulseAudio 设备: {dev['name']}")
            break

    # 如果没有找到 PulseAudio 设备，使用默认设备
    if input_device is None:
        rospy.logwarn("未找到 PulseAudio 设备，将使用默认输入设备")
        try:
            input_device = sd.default.device[0]
            dev = sd.query_devices(input_device)
            rospy.loginfo(f"使用默认设备: {dev['name']}")
        except:
            input_device = None

    # 设置实际采样率
    actual_rate = 44100  # 保守选择

    # 创建输入流
    try:
        if input_device is not None:
            rospy.loginfo(f"尝试以 {actual_rate}Hz 采样率打开音频输入流...")
            stream = sd.InputStream(
                samplerate=actual_rate,
                channels=1,
                dtype="float32",
                blocksize=CHUNK,
                device=input_device,
            )
        else:
            rospy.loginfo("尝试使用简化配置打开音频输入流...")
            stream = sd.InputStream(
                samplerate=actual_rate, channels=1, dtype="float32", blocksize=CHUNK
            )

        with stream:
            rospy.loginfo("音频流已成功启动")
            while not rospy.is_shutdown():
                try:
                    data, overflowed = stream.read(CHUNK)
                    if overflowed:
                        rospy.logwarn("音频缓冲区溢出")

                    # 转换为 numpy 数组
                    wave = data.flatten().astype(np.float32)

                    # 非阻塞写入队列
                    try:
                        audio_queue.put_nowait(wave)
                    except Full:
                        rospy.logwarn("音频队列已满，跳过当前块")
                        try:
                            audio_queue.get_nowait()
                        except Empty:
                            pass
                except Exception as e:
                    rospy.logwarn(f"音频读取错误: {str(e)}")
                    time.sleep(0.1)

    except Exception as e:
        rospy.logerr(f"无法打开音频设备: {str(e)}")
        rospy.loginfo("尝试使用备用配置...")
        try:
            # 尝试使用更小的块大小和更低的采样率
            stream = sd.InputStream(
                samplerate=16000, channels=1, dtype="float32", blocksize=512
            )
            with stream:
                rospy.loginfo("使用备用配置成功启动音频流")
                while not rospy.is_shutdown():
                    try:
                        data, overflowed = stream.read(512)
                        if overflowed:
                            rospy.logwarn("音频缓冲区溢出")

                        wave = data.flatten().astype(np.float32)
                        try:
                            audio_queue.put_nowait(wave)
                        except Full:
                            try:
                                audio_queue.get_nowait()
                            except Empty:
                                pass
                    except Exception as e2:
                        rospy.logwarn(f"备用配置读取错误: {str(e2)}")
                        time.sleep(0.1)
        except Exception as e3:
            rospy.logerr(f"备用配置也失败: {str(e3)}")


def audio_consumer():
    """音频处理线程 - 添加错误恢复"""
    rospy.loginfo("音频处理已启动")

    last_detected_freq = None
    last_publish_time = rospy.Time.now()
    consecutive_errors = 0

    while not rospy.is_shutdown() and consecutive_errors < 10:
        try:
            wave = audio_queue.get(timeout=1.0)
            detected_freq = sound_recognize(wave)

            consecutive_errors = 0  # 重置错误计数器

            if detected_freq and detected_freq in freq_to_cmd:
                current_time = rospy.Time.now()

                # 1. 频率稳定检测
                if detected_freq == last_detected_freq:
                    # 2. 限制发布频率
                    if (current_time - last_publish_time).to_sec() > 0.1:
                        cmd = freq_to_cmd[detected_freq]
                        twist = Twist()
                        twist.linear.x = cmd["x"]
                        twist.linear.y = cmd["y"]
                        cmd_pub.publish(twist)
                        rospy.loginfo(f"检测到频率: {detected_freq}Hz → 指令: {cmd}")
                        last_publish_time = current_time
                else:
                    last_detected_freq = detected_freq
            else:
                # 没有检测到有效频率时发布停止指令
                if last_detected_freq is not None:
                    twist = Twist()
                    twist.linear.x = 0
                    twist.linear.y = 0
                    cmd_pub.publish(twist)
                    rospy.loginfo("停止指令")
                    last_detected_freq = None

        except Empty:
            continue
        except Exception as e:
            consecutive_errors += 1
            rospy.logerr(f"音频处理错误 ({consecutive_errors}/10): {str(e)}")
            time.sleep(0.1)

    if consecutive_errors >= 10:
        rospy.logerr("达到最大连续错误次数，音频处理线程退出")


if __name__ == "__main__":
    try:
        # 启动生产者线程
        producer_thread = threading.Thread(target=audio_producer)
        producer_thread.daemon = True
        producer_thread.start()

        # 启动消费者线程
        consumer_thread = threading.Thread(target=audio_consumer)
        consumer_thread.daemon = True
        consumer_thread.start()

        rospy.loginfo("超声波接收节点已启动，等待指令...")

        # 保持节点运行
        rospy.spin()

    except rospy.ROSInterruptException:
        rospy.loginfo("节点已关闭")
